import pandas as pd
import os
from decimal import Decimal, ROUND_HALF_UP
from Functions import *
from Factor_Alpha101_code import *
pd.set_option('display.max_rows', 1000)  # 最多显示数据的行数
pd.set_option('expand_frame_repr', False)  # 当列太多时显示不清楚


class DataSet(object):
    #初始属性为数据存放的目录
    def __init__(self, data_dirs, index_dir, hdf_dir):
        self.data_dirs = data_dirs
        self.index_dir = index_dir
        self.hdf_dir = hdf_dir
        self.alphalist=[]
        self.retain_factors={'周期最后交易日': 'last', '代码': 'last', '简称': 'last', '是否交易': 'last', '下日_是否交易': 'last',
                             '下日_开盘涨停': 'last', '下日_是否ST': 'last', '下日_是否退市': 'last', '下日_开盘买入涨跌幅': 'last',
                             '涨跌幅(%)':'last', '开盘价(元)':'last', '收盘价(元)':'last', '成交量(股)':'last', '涨跌幅(%)':'last'}

    def generate_alpha(self, x):
        self.alphalist = x
        for i in x:
            self.retain_factors[i]='last'

# 获取指数，从指定位置读取指数数据
    def import_index_data(self):
        # 导入指数数据
        df_index = pd.read_csv(self.index_dir, parse_dates=['candle_end_time'],encoding='gb18030')
        df_index['指数涨跌幅'] = df_index['close'].pct_change()
        df_index = df_index[['candle_end_time', '指数涨跌幅']]
        df_index.dropna(subset=['指数涨跌幅'], inplace=True)

        df_index.rename(columns={'candle_end_time': '日期'}, inplace=True)
        df_index.sort_values(by=['日期'], inplace=True)
        df_index.reset_index(inplace=True, drop=True) 

        self.df_index = df_index


    # 从指定的目录下获取所有的股票csv文件的文件名，并赋值给self.stock_list属性,并遍历stock_list获取所有股票信息
    # 包括一些数据清理和处理过程
    def get_stock_code_list_in_dir(self,period_type):
        stock_list = []
        # os.walk，用于遍历文件夹中的所有文件
        for root, dirs, files in os.walk(self.data_dirs):
            if files:  # 当files不为空的时候
                for f in files:
                    if f.endswith('.CSV'):
                        stock_list.append(f[:9])
        self.stock_list = stock_list
        all_stock_data = pd.DataFrame()  # 用于存储数据
        for code in self.stock_list:
            print(code)
            # =读入股票数据
            df = pd.read_csv(self.data_dirs + '/%s.csv' % code, encoding='gbk',parse_dates=[2])
            df=df.loc[:, '代码':'均价(元)']
            # 删除2017年之前的数据
            df = df[df['日期'] >= pd.to_datetime('20190101')]
            df = df[df['日期'] <= pd.to_datetime('20210101')]
            # =计算涨跌幅
            df['涨跌幅(%)'] = df['收盘价(元)'] / df['前收盘价(元)'] - 1
            df['开盘买入涨跌幅'] = df['收盘价(元)'] / df['开盘价(元)'] - 1  # 为之后开盘买入做好准备

            # =将股票和上证指数合并，补全停牌的日期，新增数据"是否交易"、"指数涨跌幅"
            """
            原始股票数据在不交易的时候没有数据。
            将原始股票数据和指数数据合并，可以补全原始股票数据没有交易的日期。
            """
            # ===将股票数据和上证指数合并，结果已经排序
            df = pd.merge(left=df, right=self.df_index, on='日期', how='right', sort=True, indicator=True)

            # ===对开、高、收、低、前收盘价(元)价格进行补全处理
            # 用前一天的收盘价，补全收盘价的空值
            df['收盘价(元)'].fillna(method='ffill', inplace=True)
            # 用收盘价补全开盘价(元)、最高价(元)、最低价的空值
            df['开盘价(元)'].fillna(value=df['收盘价(元)'], inplace=True)
            df['最高价(元)'].fillna(value=df['收盘价(元)'], inplace=True)
            df['最低价(元)'].fillna(value=df['收盘价(元)'], inplace=True)
            # 补全前收盘价(元)
            df['前收盘价(元)'].fillna(value=df['收盘价(元)'].shift(), inplace=True)



            # ===将停盘时间的某些列，数据填补为0
            fill_0_list = ['成交量(股)', '成交金额(元)', '涨跌幅(%)', '开盘买入涨跌幅']
            df.loc[:, fill_0_list] = df[fill_0_list].fillna(value=0)

            # ===用前一天的数据，补全其余空值
            df.fillna(method='ffill', inplace=True)

            # ===去除上市之前的数据
            df = df[df['代码'].notnull()]

            # ===判断计算当天是否交易
            df['是否交易'] = 1
            df.loc[df['_merge'] == 'right_only', '是否交易'] = 0
            del df['_merge']

            df.reset_index(drop=True, inplace=True)
            # 股票与指数数据合并到此结束


            # =计算涨跌停价格
            """
            计算股票当天的涨跌停价格。在计算涨跌停价格的时候，按照严格的四舍五入。
            包含st股，但是不包含新股
            """
            # 计算涨停价格
            df['涨停价'] = df['前收盘价(元)'] * 1.1
            # 针对st进行修改
            df.loc[df['简称'].str.contains('ST'), '涨停价'] = df['前收盘价(元)'] * 1.05
            # 四舍五入
            df['涨停价'] = df['涨停价'].apply(lambda x: float(Decimal(x * 100).quantize(Decimal('1'), rounding=ROUND_HALF_UP) / 100))
            # 判断是否一字涨停
            df['一字涨停'] = False
            df.loc[df['最低价(元)'] >= df['涨停价'], '一字涨停'] = True
            # 判断是否开盘涨停
            df['开盘涨停'] = False
            df.loc[df['开盘价(元)'] >= df['涨停价'], '开盘涨停'] = True


            # =计算下个交易的相关情况
            df['下日_是否交易'] = df['是否交易'].shift(-1)
            df['下日_一字涨停'] = df['一字涨停'].shift(-1)
            df['下日_开盘涨停'] = df['开盘涨停'].shift(-1)
            df['下日_是否ST'] = df['简称'].str.contains('ST').shift(-1)
            df['下日_是否退市'] = df['简称'].str.contains('退').shift(-1)
            df['下日_开盘买入涨跌幅'] = df['开盘买入涨跌幅'].shift(-1)

            #计算因子值
            get_alpha(df, self.alphalist)


            # 将日线数据转换为其他周期的数据
            df['周期最后交易日'] = df['日期']
            df.set_index('日期', inplace=True)
            # print(df.columns)
            period_df = df.resample(rule=period_type).agg(self.retain_factors)

            # 计算必须额外数据
            period_df['交易天数'] = df['是否交易'].resample(period_type).sum()
            period_df['市场交易天数'] = df['代码'].resample(period_type).size()
            period_df = period_df[period_df['市场交易天数'] > 0]  # 有的时候整个周期不交易（例如春节、国庆假期），需要将这一周期删除

            # 计算其他因子

            # 计算周期资金曲线
            period_df['每天涨跌幅'] = df['涨跌幅(%)'].resample(period_type).apply(lambda x: list(x))
            
            # 重新设定index
            period_df.reset_index(inplace=True)
            period_df['交易日期'] = period_df['周期最后交易日']
            del period_df['周期最后交易日']

            df = period_df



            df.drop([0], axis=0, inplace=True)  # 删除第一行数据
            # 计算下周期每天涨幅
            df['下周期每天涨跌幅'] = df['每天涨跌幅'].shift(-1)

            df = df[df['简称'].str.contains('ST') == False]
            # 删除月末有退市风险的周期数
            df = df[df['简称'].str.contains('退') == False]
            # 删除月末不交易的周期数
            df = df[df['是否交易'] == 1]
            # 删除交易天数过少的周期数

            # 合并数据，存储到all_stock_data中
            all_stock_data = all_stock_data.append(df, ignore_index=True)

        self.all_stock_data = all_stock_data


    # 对数据进行排序，重新设定索引,并存入hdf文件
    def sort_reindex_store(self, sortby:list):
        '''
        对数据进行排序，重新设定索引,并存入hdf文件
        param: sortby,column组成的数组，排序的关键字
        '''
        # 将数据存入数据库之前，先排序、reset_index
        self.all_stock_data.sort_values(sortby, inplace=True)
        self.all_stock_data.reset_index(inplace=True, drop=True)

        # 将数据存储到hdf文件
        # self.all_stock_data.to_hdf(self.hdf_dir + r'/all_stock_data_d'+r'.h5', 'df', format = 'fixed', mode='w')
        print("pd2hdf5:Completed:" + self.hdf_dir)


